"""
@Time       :   2024/06/13
@Author     :   Yi Junquan
@Version    :   1.0
@Contact    :   2696974822@qq.com
@Software   :   VsCode
"""

from util.types import NRange

class MetricBase:
    """
    基础度量类，包含共享的属性和方法。
    """
    def __init__(self, anonymized_data, qi_index) -> None:
        self.anonymized_data = anonymized_data
        self.qi_index = qi_index
        self.qi_number = len(qi_index)
        # 等价类统计
        self.eq_class_count = {}

    def calculate_eq_class_count(self):
        """
        计算等价类的数量
        """
        self.eq_class_count = {}
        for record in self.anonymized_data:
            qi_values = tuple(record[qi_id] for qi_id in self.qi_index)
            if qi_values not in self.eq_class_count:
                self.eq_class_count[qi_values] = 0
            self.eq_class_count[qi_values] += 1

class NCP(MetricBase):
    """
    normalized certainty penalty的实现
    """
    def __init__(self, anonymized_data, qi_index, tax_trees) -> None:
        super().__init__(anonymized_data, qi_index)
        self.tax_trees = tax_trees
        self.get_cat_range_leaves()

    def get_cat_range_leaves(self):
        self.qi_range = []
        self.cat_flag = []
        self.qi_leaves = []
        for t in self.tax_trees:
            if not isinstance(t, NRange):
                self.cat_flag.append(True)
            else:
                self.cat_flag.append(False)
        for i in range(self.qi_number):
            self.qi_leaves.append({})
            if self.cat_flag[i]:
                self.qi_range.append(len(self.tax_trees[i]['*']))        
            else:
                self.qi_range.append(self.tax_trees[i].range)

            if self.cat_flag[i]:
                for key in self.tax_trees[i].keys():
                    self.qi_leaves[i][key] = len(self.tax_trees[i][key])

    def get_evaluation_score(self):
        """
        计算NCP分数
        """
        ncp = 0.0 
        for record in self.anonymized_data:
            ncp_single = 0.0
            for idx, qi_id in enumerate(self.qi_index):
                value = record[qi_id]
                if not self.cat_flag[idx]: #数值类型
                    low, high = value.split('-')
                    ncp_single += ((float(high) - float(low)) / self.qi_range[idx])
                else:
                    num_leaves = self.qi_leaves[idx][value]
                    ncp_single += (num_leaves * 1.0 / self.qi_range[idx])
            ncp += ncp_single
        ncp /= self.qi_number
        ncp /= len(self.anonymized_data)
        return ncp


class CAVG(MetricBase):
    """
    normalized average equivalence class size的实现   
    """
    def __init__(self, anonymized_data, qi_index, k) -> None:
        super().__init__(anonymized_data, qi_index)
        self.k = k

    def get_evaluation_score(self):
        self.calculate_eq_class_count()
        num_records = len(self.anonymized_data)
        num_eqs = len(self.eq_class_count)
        return num_records * 1.0 / (num_eqs * self.k)


class DM(MetricBase):
    """discernability model的实现
    """
    def __init__(self, anonymized_data, qi_index, k) -> None:
        super().__init__(anonymized_data, qi_index)
        self.k = k
        self.num_records = len(anonymized_data)

    def get_evaluation_score(self):
        self.calculate_eq_class_count()
        dm = 0
        for eq_class_count in self.eq_class_count.values():
            if eq_class_count < self.k:
                dm += (eq_class_count * self.num_records)
            else:
                dm += (eq_class_count * eq_class_count)       
        return dm
